-------------------------------------------------------------------------------
-- Copas - Coroutine Oriented Portable Asynchronous Services
-- 
-- Offers a dispatcher and socket operations based on coroutines.
-- Usage:
--    copas.addserver(server, handler, timeout)
--    copas.addthread(thread, ...) Create a new coroutine thread and run it with args
--    copas.loop(timeout) - listens infinetely
--    copas.step(timeout) - executes one listening step
--    copas.receive(pattern or number) - receives data from a socket
--    copas.settimeout(client, time) if time=0 copas.receive(bufferSize) - receives partial data from a socket were data<=bufferSize
--    copas.send  - sends data through a socket
--    copas.wrap  - wraps a LuaSocket socket with Copas methods
--    copas.connect - blocks only the thread until connection completes
--    copas.flush - *deprecated* do nothing
--
-- Authors: Andre Carregal and Javier Guerra
-- Contributors: Diego Nehab, Mike Pall, David Burgess, Leonardo Godinho
--               and Thomas Harning Jr.
--
-- Copyright 2007 - Kepler Project (www.keplerproject.org)
--
-- $Id: copas.lua,v 1.26 2007/06/20 21:57:16 carregal Exp $
--
-- Partial copyright 2007 gary ng<linux@garyng.com>
--
-- the async IO part has been factored out so this now is a pure dispatcher
-- using select(). the coroutine yield also returns string rather than internal
-- buffer making it easier to adopt from the caller. Currently there is three 
-- states : reading, writing, sleeping
-------------------------------------------------------------------------------
local socket = require "socket"
local asyncio = require 'asyncio'
require 'coxpcall'
local time = os.time
local duration = os.difftime
local tinsert, tremove = table.insert, table.remove

WATCH_DOG_TIMEOUT = 120

-- Redefines LuaSocket functions with coroutine safe versions
-- (this allows the use of socket.http from within copas)
local function statusHandler(status, ...)
  if status then return ... end
  return nil, ...
end

function socket.protect(func)
  --if true then return function(...) return func(...) end end
  return function (...)
		local ret = {copcall(func, ...)}
		local status = table.remove(ret,1)
		if status then
			return unpack(ret)
		end
		return nil, unpack(ret)
	end
end

function socket.newtry(finalizer)
	return function (...)
    local arg={...}
		local status = arg[1]or false
		if (status==false)then
			table.remove(arg,1)
			local ret = {pcall(finalizer, unpack(arg) ) }
			error(arg[1], 0) --this would terminate a running thread(coroutine)
		end
		return unpack(arg)
	end
end
-- end of LuaSocket redefinitions


module ("copas1", package.seeall)

-- Meta information is public even if begining with an "_"
_COPYRIGHT   = "Copyright (C) 2004-2007 Kepler Project"
_DESCRIPTION = "Coroutine Oriented Portable Asynchronous Services"
_VERSION     = "Copas 1.1.1"

-------------------------------------------------------------------------------
-- Simple set implementation based on LuaSocket's tinyirc.lua example
-- adds a FIFO queue for each value in the set
-------------------------------------------------------------------------------
local function newset()
    local reverse = {}
    local set = {}
    local q = {}
    setmetatable(set, { __index = {
        __index = function(set, k) return reverse[k] end,

        insert = function(set, value)
            if not reverse[value] then
                set[#set + 1] = value
                reverse[value] = #set
            end
        end,

        remove = function(set, value)
            local index = reverse[value]
            if index then
                reverse[value] = nil
                local top = set[#set]
                set[#set] = nil
                if top ~= value then
                    reverse[top] = index
                    set[index] = top
                end
            end
        end,
		
		push = function (set, key, itm)
			local qKey = q[key]
			if qKey == nil then
				q[key] = {itm}
			else
				qKey[#qKey + 1] = itm
			end
		end,
		
        pop = function (set, key)
          local t = q[key]
          if t ~= nil then
            local ret = table.remove (t, 1)
            if t[1] == nil then
              q[key] = nil
            end
            return ret
          end
        end,
    }})
    return set
end

local _servers = newset() -- servers being handled

local _selectable = {reading = newset(), writing=newset(), sleeping=newset()}
local _threads = {}
local _dangling = {}
local _killing = {}

-------------------------------------------------------------------------------
-- Coroutine based socket I/O functions.
-------------------------------------------------------------------------------
--
--an ojbect mapping if the send/receive method is used
--instead of the wrapper

local obj_map=setmetatable({},{__mode="kv"}) 

function wrap (skt, buffered_write)
	local s =  asyncio.wrap(skt, buffered_write)
    return s
end

-- reads a pattern from a client and yields to the reading set on timeouts
function receive(obj, pattern, ...)
    local state = obj_map[obj]
    if not state then 
        state = wrap(obj)
        obj_map[obj] = state
    end

    return state:receive(pattern, ...)
end

-- same as above but with special treatment when reading chunks,
-- unblocks on any data received.
function receivePartial(obj, pattern, ...)
    local state = obj_map[obj]
    if not state then 
        state = wrap(obj)
        obj_map[obj] = state
    end

    state:settimeout(0)
    return state:receive(pattern, ...)
end

-- sends data to a client. The operation is buffered and
-- yields to the writing set on timeouts
function send(obj, data, ...)
    local state = obj_map[obj]
    if not state then 
        state = wrap(obj)
        obj_map[obj] = state
    end

    return state:send(data, ...)
end

-- waits until connection is completed
function connect(skt, host, port)
  if non_block_connect == false  then 
    skt:settimeout()
    return skt:connect(host, port)
  end
  local obj_key = skt
  local state = obj_map[obj_key]
  if not state then 
      state = wrap(skt)
      obj_map[obj_key] = state
  end
  return state:connect(host, port)
end

-- flushes a client write buffer (deprecated)
function flush(client)
end

--------------------------------------------------
-- Error handling
--------------------------------------------------

local _errhandlers = {}   -- error handler per coroutine

function setErrorHandler (err)
	local co = coroutine.running()
	if co then
		_errhandlers [co] = err
	end
end

local function _deferror (msg, co, skt)
	print (msg, co, skt)
end

-------------------------------------------------------------------------------
-- Thread handling
-------------------------------------------------------------------------------

local function _doTick (new, co, skt, ...)
	if not co then return end
	local ok, res, new_q, ev = coroutine.resume(co, skt, ...)
	if ok and res and new_q and not _killing[co] then
    if skt then
      local pending = _threads[skt]
      if pending then
        local active_threads = #pending.reading + #pending.writing + #pending.sleeping
        if active_threads < 1 then _threads[skt] = nil end
      end
    end
    local dead = co and _dangling[co]
    if not dead or os.time() - dead < WATCH_DOG_TIMEOUT*2 then
      _selectable[new_q]:insert(res)
      local pending = _threads[res]
      if not pending then
          pending = {reading={},writing={},sleeping={}}
         _threads[res] = pending
      end
      if new then 
          tinsert(pending[new_q], {co, res, time(), ev})
      else 
          tinsert(pending[new_q],1, {co, res, time(), ev}) 
      end
      dead = nil
    else
      --_killing[co] = os.time()
      --_dangling[co] = nil
      --res:close()
      return _doTick(false, co, res, true)
    end
	else
    local killing = _killing[co]
    if co then 
      if coroutine.status(co) ~= "dead" then
        --_killing[co] = nil
        _dangling[co] = nil
        return _doTick(false, co, res, true)
      end
      _killing[co] = nil
      _dangling[co] = nil 
    end
		if skt then
        if not new then
            local pending = _threads[skt]
            if pending then
                local active_threads = #pending.reading + #pending.writing + #pending.sleeping
                if active_threads > 0 then
                    --print("active threads", active_threads)
                    return
                end
                _threads[skt] = nil 
            end
        end
        --skt:close() 
    end
    local err_handler = _errhandlers [co] or _deferror
		_errhandlers [co] = nil
		if not ok then 
      local s,skip = pcall (err_handler, res, co, skt) 
      --if not skip then error(res, new_q, ev) end
    end
	end
end

-- accepts a connection on socket input
local function _accept(input, handler)
  local client
  if input.receivefrom then client = input
  else client = input:accept() end
	if client then 
		client:settimeout(0)
		local co = coroutine.create(handler)
		_doTick (true, co, client)
		--_reading:insert(client)
	end
	return client
end

-- handle threads on a queue
local function _tickRead (skt)
    local pending = _threads[skt]
    if pending then
        if pending.reading[1] then
            local co, v, t = unpack(tremove(pending.reading, 1))
            _doTick (false, co, v)
        elseif pending.writing[1] then
            --print("readable but waiting in write")
        end
    else
        --print('where is the read socket', skt)
    end
end

local function _tickWrite (skt)
    local pending = _threads[skt]
    if pending then
        if pending.writing[1] then
            local co, v, t = unpack(tremove(pending.writing,1))
            _doTick (false, co, v)
        elseif pending.reading[1] then
            --print("writable but waiting in read")
        end
    else
        --print('where is the write socket', skt)
    end
end

local function _tickSleep (skt)
    local pending = _threads[skt]
    _selectable['sleeping']:remove(skt)
    if pending then
        if pending.sleeping[1] then
            local co, v, t = unpack(tremove(pending.sleeping,1))
            pending[skt] = nil
            _doTick (false, co, v)
        end
    end
end

-------------------------------------------------------------------------------
-- Adds a server/handler pair to Copas dispatcher
-------------------------------------------------------------------------------
function addserver(server, handler, timeout)
  server:settimeout(timeout or 0.1)
  _servers[server] = handler
  _selectable['reading']:insert(server)
end

function delserver(server)
  _selectable['reading']:remove(server)
  _selectable['writing']:remove(server)
  _selectable['sleeping']:remove(server)
  _servers[server] = nil
end

-------------------------------------------------------------------------------
-- Adds an new courotine thread to Copas dispatcher
-------------------------------------------------------------------------------
function addthread(thread, ...)
	local co = coroutine.create(thread)
	_doTick (true, co, nil, ...)
    return co
end

-------------------------------------------------------------------------------
-- tasks registering
-------------------------------------------------------------------------------

local _tasks = {}

local function addtaskRead (tsk)
	-- lets tasks call the default _tick()
	tsk.def_tick = _tickRead
	
	_tasks [tsk] = true
end

local function addtaskWrite (tsk)
	-- lets tasks call the default _tick()
	tsk.def_tick = _tickWrite
	
	_tasks [tsk] = true
end

local function addtaskSleep (tsk)
	-- lets tasks call the default _tick()
	tsk.def_tick = _tickSleep
	
	_tasks [tsk] = true
end

local function tasks ()
	return next, _tasks
end

-------------------------------------------------------------------------------
-- main tasks: manage readable and writable socket sets
-------------------------------------------------------------------------------
-- a task to check ready to read events
local _readable_t = {
  events = function(self)
  	local i = 0
  	return function ()
  		i = i + 1
  		return self._evs [i]
  	end
  end,
  
  tick = function (self, input)
  	local handler = _servers[input]
  	if handler then
  		input = _accept(input, handler)
  	else
  		_selectable['reading']:remove (input)
  		self.def_tick (input)
  	end
  end
}

addtaskRead (_readable_t)


-- a task to check ready to write events
local _writable_t = {
  events = function (self)
  	local i = 0
  	return function ()
  		i = i+1
  		return self._evs [i]
  	end
  end,
  
  tick = function (self, output)
  	_selectable['writing']:remove (output)
  	self.def_tick (output)
  end
}

addtaskWrite (_writable_t)

-- a task to check ready to write events
local _wakeable_t = {
  events = function (self)
  	local i = 0
  	return function ()
  		i = i+1
  		return self._evs [i]
  	end
  end,
  
  tick = function (self, output)
  	self.def_tick (output)
  end
}

addtaskSleep (_wakeable_t)
--
local timers={}
local new_timer = false

local timer_mt = {
    close = function() return end,
    getfd = function(self) return self end,
}; timer_mt.__index = timer_mt

function sleep(timeout)
    local t = setmetatable({os.time() + timeout},timer_mt)

    timers[#timers + 1] = t
    new_timer = true
    coroutine.yield(t, "sleeping")
end

local last_cleansing = time()
local function _select (timeout)
    local now = time()

	_readable_t._evs, _writable_t._evs, err = socket.select(_selectable['reading'], _selectable['writing'], timeout)


    local s_evs = {}
    _wakeable_t._evs = s_evs

    if new_timer then 
        new_timer = false
        table.sort(timers,function(a,b) return a[1] < b[1] end) 
    end

    for _,v in ipairs(timers) do
        if v[1] < now then s_evs[#s_evs + 1] = v end
    end
    

	local r_evs, w_evs  = _readable_t._evs, _writable_t._evs 

    if duration(now, last_cleansing) > WATCH_DOG_TIMEOUT then
        last_cleansing = now
        for _,job in pairs(_threads) do
            local r = job.reading[1]
            local w = job.writing[1]
            if r then
                local co, skt, t = unpack(r)
                if co and duration(now, t) > WATCH_DOG_TIMEOUT and not r_evs[skt] then
                  if not _dangling[co] then _dangling[co] = now end
                   r_evs[#r_evs + 1] = skt 
                   r_evs[skt] = #r_evs
                end
            end

            if w then
                local co, skt, t = unpack(w)
                if co and duration(now, t) > WATCH_DOG_TIMEOUT and not w_evs[skt] then
                  if not _dangling[co] then _dangling[co] = now end
                   w_evs[#w_evs + 1] = skt 
                   w_evs[skt] = #w_evs
                end
            end
        end
    end

    if err == "timeout" and #r_evs + #w_evs + #s_evs > 0 then 
      return nil
    else return err end
end


-------------------------------------------------------------------------------
-- Dispatcher loop step.
-- Listen to client requests and handles them
-------------------------------------------------------------------------------
function step(timeout)
  local err = _select (timeout)
  if err == "timeout" then 
    return err
  end

	if err then
		error(err)
	end
	for tsk in tasks() do
		for ev in tsk:events () do
			tsk:tick (ev)
		end
	end
end

function state()
  local i = 0
  for k,v in pairs(_threads) do 
    i = i + 1
  end
  return i
end
-------------------------------------------------------------------------------
-- Dispatcher endless loop.
-- Listen to client requests and handles them forever
-------------------------------------------------------------------------------
local loop_done

function loopexit()
  loop_done = true
end

function loop(timeout)
	while not loop_done do
		step(timeout or 0.1)
	end
end
